package com.peng.leetcode.trie;

import java.util.HashMap;
import java.util.Map;

/**
 * TrieTree
 *
 * @author: lupeng6
 * @create: 2021/1/29 16:17
 */
public class TrieTree {

    public static void main(String[] args) {
        TrieTree trie = new TrieTree();

        trie.insert("apple");
        System.out.println(trie.search("apple"));
        System.out.println(trie.search("app"));
        System.out.println(trie.startsWith("app"));
        trie.insert("app");
        System.out.println(trie.search("app"));
    }


//    /** Initialize your data structure here. */
//    public Trie() {
//
//    }

    private Node root = new Node(null);

    /**
     * Inserts a word into the trie.
     */
    public void insert(String word) {
        Map<Character, Node> links = root.links;
        char c = word.charAt(0);
        links.putIfAbsent(c, new Node(c));
        links.get(c).insert(word);
    }

    /**
     * Returns if the word is in the trie.
     */
    public boolean search(String word) {
        char c = word.charAt(0);
        Node node = root.links.get(c);
        if (node == null) {
            return false;
        }
        return node.search(word);
    }

    /**
     * Returns if there is any word in the trie that starts with the given prefix.
     */
    public boolean startsWith(String prefix) {
        char c = prefix.charAt(0);
        Node node = root.links.get(c);
        if (node == null) {
            return false;
        }
        return node.startsWith(prefix);
    }

    static class Node {
        Character value;
        Map<Character, Node> links = new HashMap<>();
        boolean isEnd = false;

        public Node(Character value) {
            this.value = value;
        }

        public void insert(String word) {
            if (word.length() <= 1) {
                isEnd = true;
            } else {
                char c = word.charAt(1);
                links.putIfAbsent(c, new Node(c));
                links.get(c).insert(word.substring(1));
            }
        }

        public boolean search(String word) {
            Node node = searchLastNode(word);
            if (node == null) {
                return false;
            }
            return node.isEnd;
        }

        private Node searchLastNode(String word) {
            if (word.length() == 1) {
                return this;
            }
            char c = word.charAt(1);
            Node node = links.get(c);
            if (node == null) {
                return null;
            }
            return node.searchLastNode(word.substring(1));
        }

        public boolean startsWith(String prefix) {
            Node node = searchLastNode(prefix);
            return node != null;
        }
    }

}
